import os
import torch
import torch.nn as nn
import torch.optim as optim
import gc
from data_processing import load_and_process_data, create_data_loaders, create_sample_data
from models import TextCNN, train_word2vec
from train import train_model, load_model
from predict import predict, save_predictions

def main():
    os.makedirs('./models', exist_ok=True)
    os.makedirs('./output', exist_ok=True)
    
    TRAIN_PATH = './data/train_set.csv'
    TEST_PATH = './data/test_a.csv'
    MODEL_PATH = './models/textcnn_model.pth'
    OUTPUT_PATH = './output/submit_textcnn.csv'
    
    if not os.path.exists(TRAIN_PATH) or not os.path.exists(TEST_PATH):
        create_sample_data()
    print("正在加载和处理数据...")
    train_texts, train_labels, test_texts = load_and_process_data(
        TRAIN_PATH, TEST_PATH
    )
    
    train_loader, val_loader, test_loader = create_data_loaders(
        train_texts, train_labels, test_texts, batch_size=64
    )
    
    print("正在训练Word2Vec模型...")
    pretrained_embeddings = train_word2vec(train_texts)
    gc.collect()
    
    print("正在初始化TextCNN模型...")
    model = TextCNN(
        vocab_size=100000,
        embed_dim=128,
        num_classes=14,
        pretrained_embeddings=pretrained_embeddings
    )
    
    print("开始训练模型...")
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    
    model = train_model(
        model, train_loader, val_loader,
        criterion, optimizer,
        num_epochs=5,
        save_path=MODEL_PATH
    )
    
    print("加载最佳模型...")
    model = load_model(model, MODEL_PATH)
    
    print("正在预测测试集...")
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = model.to(device)
    predictions = predict(model, test_loader, device)
    
    print("正在保存预测结果...")
    save_predictions(predictions, OUTPUT_PATH)
    
    print("预测完成")

if __name__ == "__main__":
    main()
